# 将元数据随机打乱，并选取10%作为测试集（90%训练+验证，10%测试）
# import splitfolders
# input_folder = "C:\Temp\TF2.4 training\dataset_root"
# output_folder = "C:\Temp\TF2.4 training\splited dataset"
#
# splitfolders.ratio(
#     input_folder,
#     output=output_folder,
#     seed=42,
#     ratio=(0.9, 0.1),
#     group_prefix=None
# )

# 遍历文件夹，生成标签文件
import os

# 1
# with open('labels.txt', 'w') as f:
#     for label, class_name in enumerate(os.listdir(datasets_dir)):
#         for img_path in os.listdir(f'dataset/{class_name}'):
#             f.write(f'dataset/{class_name}/{img_path}, {label}\n')

# 2
# def generate_img_labels(folder_path, output_path, label):
#     files = os.listdir(folder_path)
#     files.sort()
#     with open(output_path, 'w') as f:
#         for file in files:
#             if not file.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp')):
#                 continue
#             f.write(f"{file} {label}\n")
#
# generate_img_labels(dataset_28, output_label_28, 0)
# generate_img_labels(dataset_36, output_label_36, 1)

# 3
dataset_train = 'C:\Temp\TF2.4 training\splited dataset\\train'
dataset_test = 'C:\Temp\TF2.4 training\splited dataset\\test'
output_label_train = 'C:\Temp\TF2.4 training\splited dataset\labels-train.txt'
output_label_test = 'C:\Temp\TF2.4 training\splited dataset\labels-test.txt'

def generate_img_labels(folder_path, output_path):
    files = os.listdir(folder_path)
    files.sort()
    with open(output_path, 'w') as f:
        for file in files:
            if not file.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp')):
                continue

            if file[:2] == '28':
                f.write(f"{file} {0}\n")
            else:
                f.write(f"{file} {1}\n")

generate_img_labels(dataset_train, output_label_train)
generate_img_labels(dataset_test, output_label_test)